import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_blobs

data, true_labels = make_blobs(centers = 5, random_state = 20, cluster_std = 1) # data是坐标集合，ture_labels是正确标签
plt.scatter(data[:, 0], data[:, 1])
plt.show()


class my_Kmeans:
    def __init__(self, k):
        self.k = k

    
    def calc_distance(self, vec1, vec2):
        return np.sqrt(np.sum(np.power(vec1 - vec2, 2)))

    def fit(self, data):
        numSamples, dim = data.shape 
        
        self.centers_idx = np.random.choice(numSamples, self.k, replace=False) 
        self.centers = data[self.centers_idx].astype(np.float32) 

        
        ClusterAssment = np.zeros((numSamples, 2))
        ClusterChanged = True

        while ClusterChanged:
            ClusterChanged = False
            
            for i in range(numSamples):
                mindist = self.calc_distance(data[i], self.centers[0])
                label = 0
                for j in range(1, self.k):
                    distance = self.calc_distance(data[i], self.centers[j])
                    
                    if distance < mindist:
                        mindist = distance
                        label = j

                

                
                if ClusterAssment[i, 0] != label:
                    ClusterChanged = True
                    ClusterAssment[i, :] = label, mindist ** 2

           

            
            for j in range(self.k):
                
                pointsInCluster = data[ClusterAssment[:, 0] == j]
                
                self.centers[j, :] = (np.mean(pointsInCluster, axis=0).tolist()) # 更新聚类中心的位置

    def predict(self, data):
        numSamples, dim = data.shape
        ClusterAssment = np.zeros((numSamples, 2))
        for i in range(numSamples):
            mindist = self.calc_distance(data[i], self.centers[0])
            label_pred = 0
            for j in range(1, self.k):
                distance = self.calc_distance(data[i], self.centers[j])
                if distance < mindist:
                    mindist = distance
                    label_pred = j

            ClusterAssment[i, :] = label_pred, mindist ** 2

        return self.centers, ClusterAssment

model = my_Kmeans(k=5)
model.fit(data)
centers, ClusterAssment = model.predict(data)

plt.figure(figsize=(6, 5))
plt.scatter(data[:, 0], data[:, 1], c = ClusterAssment[:, 0])
plt.scatter(centers[:, 0], centers[:, 1], marker = '*', color = 'r', s = 100) # 绘制类中心点
plt.show()


